Skip to content

feat(desktop): NIP-AM agent-usage backend — P2 emission/transport/archive + P4a aggregation/D6 - #4000

Open
wpfleger96 wants to merge 2 commits into
mainfrom
duncan/agent-usage-backend
Open

feat(desktop): NIP-AM agent-usage backend — P2 emission/transport/archive + P4a aggregation/D6#4000
wpfleger96 wants to merge 2 commits into
mainfrom
duncan/agent-usage-backend

Conversation

@wpfleger96

@wpfleger96 wpfleger96 commented Jul 31, 2026

Copy link
Copy Markdown
Member

What

Implements Phases 2 and 4a of the Usage v2 plan (plan events d0268cd0/0e95b035), extending the archive backend to emit, transport, archive, and aggregate both cache categories and billing identity fail-closed.

P2 — emission, transport, archive

Tri-state accumulators (Unseen/Exact/Unknown) for cache-read and cache-write in buzz-agent turn and session state. Absent field = Unknown (never zero) through the full pipeline. No unwrap_or(0) on the cache path. Both cache folds are gated on usage-bearing responses (same gate as the total-state and identity folds) — a response with no usage at all must not poison either accumulator.

Overflow-aware input token parsing and accumulation — closed end-to-end from parse through wire to ACP:

  • sum_usage() returns SumUsageResult (Exact(u64) | Overflow) — checked arithmetic, never clamps. anthropic_input_tokens() returns Option<SumUsageResult> since it sums three fields (input_tokens + cache_read_input_tokens + cache_creation_input_tokens) that can collectively overflow. Single-field callers (prompt_tokens, completion_tokens, etc.) convert via .into_exact() — their single-field sums cannot overflow.
  • LlmResponse.input_tokens_overflowed: bool propagates the parse-layer signal into the run loop. When set, input_tokens is None (clamped value discarded), the context-gate baseline (last_request_input_tokens) is frozen at its prior reading, and turn_input_tokens is poisoned to TurnIOState::Poisoned before any emission — including mid-turn emit_usage_update calls. A dedicated enum on LlmResponse.input_tokens would ripple into ~20 existing test assertions on r.input_tokens == Some(...); the bool flag confines the change to the two call sites that check it.
  • TurnIOState (Unseen/Exact/Poisoned) for input and output: per-round fold uses checked_add; overflow poisons permanently at turn and session level, no healing. Absence does not poison (pass-2-cleared contract unchanged). Wire emission omits accumulatedInputTokens/accumulatedOutputTokens when poisoned — never null, never u64::MAX. ACP treats absent = publisher-poisoned: delta_reliable: false, null turn fields, null cumulative for that category; session cumulative stays unknown for all subsequent turns once poisoned.

Conditional wire emission for accumulatedCachedInputTokens and new accumulatedCacheWriteTokens: fields are omitted when the cumulative is Unseen or Unknown. ACP _goose/unstable/session/update contract documented next to the payload with tests for all absence/zero variants.

PricingIdentity stamping (publisher-side):

  • pricing_authority(): canonical parsed-URL endpoint comparison against the official allowlist — HTTPS only, exact allowlisted host (lookalike-safe), default port (omitted or explicit :443), required API base path, rejects userinfo/query/fragment/path-prefix lookalikes.
  • Model: the actually-requested request_model after mesh/auto resolution (not effective_model_str).
  • Turn discipline: identity retained only while ALL usage in the current turn carries one identical proven identity; any mismatch, unproven-usage-bearing response, or unpaired cumulative snapshot poisons to absent; a later matching notification does not heal a mixed turn.

ACP UsageTracker identity fold: per-in-flight-turn tri-state identity accumulator replacing last-update-wins. Any absent identity on a token-advancing notification or exact mismatch poisons to absent; poison survives later updates; reset in begin_turn()/take(); reset also when a request fails (baseline cleared so preflight gate cannot stay frozen sub-threshold on retries).

M3 migration: adds turn_cache_write_tokens, cumulative_cache_write_tokens, pricing_authority, pricing_model, pricing_cache_class to agent_metric_index. Additive, idempotent, guarded per-column by marker. M2 migration also guarded per-column (turn and cumulative cache-read columns checked and added independently; marker commits only after both are present). Fresh-DB schema includes all columns.

First-turn baselines: seed_zero_baseline seeds last_input: Some(0), last_output: Some(0), last_cached_input: Some(0), last_cache_write: Some(0), and last_total: Some(0) — all have the known-zero-at-spawn argument. Absent fields from incoming snapshots still produce unknown (tri-state unchanged). Sessions buzz-acp did not spawn (no seed) remain fail-closed on turn one.

ReportedUsage TS mirror: cacheReadTokens, cacheWriteTokens, freshInputTokens added to tauriArchive.ts as UsageField members, field-for-field with the Rust struct.

P4a — aggregation layer

Extended S-1 ladder to cache-read and cache-write via the same ladder_token path as the existing token fields.

freshInputTokens derivation: checked arithmetic, fail-closed — absent cache fields produce Unknown (not zero), overflow and cacheRead+cacheWrite > input both produce incomplete: true. Aggregated as a UsageField.

D6 comparator: sort_value() = provider total when known, else input+output when both known, else None (unknown-last). Replaces the prior total-only comparator for both agent-level and model-level sort. Ships a pinned test vector that the TS render layer (P5) must match.

Test coverage

  • buzz-agent: 440 lib + 15 integration (golden_transcripts) — includes 13 new cache_total_state_tests; 14 new turn_io_state_tests; 3 new sum_usage_* tests (exact single-field, exact two-field, overflow signals correctly); 3 new parse_anthropic_* tests (overflow flag set + value cleared, normal sum no flag, absent usage no flag); end-to-end golden transcript drives real subprocess with Anthropic-shaped input_tokens: u64::MAX, cache_read: 1 response and asserts accumulatedInputTokens absent from the emitted usage_update — no logic duplication; 3 wire pin tests; 4 fold_pricing_identity_* tests; pricing_authority() explicit-:443 acceptance
  • buzz-acp: 700 tests (691 lib + 9 integration) — 4 new usage tests (absent input → unreliable+null; absent output → unreliable+null; goose-shaped both present unchanged; poison mid-session); 3 ACP behavior tests; 7 pool lifecycle tests
  • Desktop (Rust): 2259+ tests — 14 new P4a pinned tests; 2 M3 round-trip tests; 1 serde key-shape test; 2 M2 partial-schema migration tests; first-turn cache round-trip test

Related PRs

@wpfleger96
wpfleger96 requested a review from a team as a code owner July 31, 2026 20:05
@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch from 1613f34 to 12f1f29 Compare July 31, 2026 20:53
@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch 2 times, most recently from 635b030 to 39e1a91 Compare August 3, 2026 18:56
@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch from 39e1a91 to 66a09a8 Compare August 4, 2026 01:03
@wpfleger96 wpfleger96 changed the title feat(desktop): NIP-AM agent-usage archive backend + cache-read schema feat(desktop): NIP-AM agent-usage backend — P2 emission/transport/archive + P4a aggregation/D6 Aug 4, 2026
@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch 3 times, most recently from 97baa8c to 5976df7 Compare August 4, 2026 02:53
Comment thread desktop/src-tauri/src/archive/store_migrations.rs Outdated
Comment thread crates/buzz-acp/src/usage.rs Outdated
@wesbillman

Copy link
Copy Markdown
Collaborator

Commenting on Wes’s behalf. [P2] Input/output overflow is silently published as an exact saturated count.

Per-round accumulation and turn-to-session emission use saturating_add (crates/buzz-agent/src/agent.rs:239-243, 375-382). Two valid provider reports of u64::MAX and 1 therefore produce u64::MAX; ACP treats that as a reliable delta and the archive reports it as exact. This conflicts with the checked/unknown behavior implemented for total and cache counters in this PR.

Please use checked tri-state accumulation (or otherwise poison reliability) for input/output, with both within-turn and turn-to-session overflow tests.

@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch from 5976df7 to b564ecd Compare August 4, 2026 16:15

@wesbillman wesbillman left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commenting on Wes’s behalf.

Requesting changes because the remaining overflow blocker is still present at head b564ecd99eee8d3700c422d87c43787447add96a.

The updated commit fixes the two inline findings (independent M2 column repair and known-zero first-turn cache baselines), but input/output accumulation still uses saturating_add both within a turn and when merging a completed turn into the session (crates/buzz-agent/src/agent.rs:239-242, 375-381; crates/buzz-agent/src/lib.rs:762-767). Overflow therefore remains silently serialized as an exact u64::MAX cumulative count, unlike the checked/unknown behavior now used for total and cache counters.

Please use overflow-aware tri-state/reliability handling for input and output, with tests for both within-turn and turn-to-session overflow. This is the unresolved P2 reported in #4000 (comment).

@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch from b564ecd to 45ae876 Compare August 5, 2026 21:25
…) + P4a (aggregation/D6)

P2 — cache-write + pricing-identity emission, transport, and archive:

- buzz-agent/types: CacheTotalState enum (Unseen/Exact/Unknown) for per-turn
  and per-session cache-read and cache-write accumulators. Absent field on any
  usage-bearing response permanently poisons the accumulator (Unknown); explicit
  zero is distinct from absence. No unwrap_or(0) on the cache path. 13 pinned
  tests in cache_total_state_tests: fold semantics (Some/absent/overflow/no-heal),
  exact_value() (pins wire omission), and merge_session() propagation.
- buzz-agent/agent: fold() on every usage-bearing response for both cache
  categories; gated identically to the total-state and identity folds. A response
  with no usage at all must not poison either cache accumulator. merge_session()
  at turn boundary. turn_cached_input_tokens and turn_cache_write_tokens use
  CacheTotalState throughout.
- buzz-agent/lib: session-cumulative cache fields also use CacheTotalState;
  merge_session propagates Unknown when any turn was poisoned.
- buzz-agent/wire: conditional wire emission for accumulatedCachedInputTokens
  and accumulatedCacheWriteTokens; skip when cumulative is Unseen or Unknown.
  ACP contract documented next to the payload.
- buzz-agent/config: pricing_authority() rewritten using url crate — canonical
  parsed-URL comparison, HTTPS only, exact allowlisted host, accepts both
  omitted and explicit :443 as equivalent default port, required path where
  applicable, rejects lookalikes, userinfo, query, fragment.
- buzz-agent/agent: PricingIdentity turn discipline — fold_pricing_identity()
  pure helper; identity retained only while all usage-bearing rounds carry one
  identical proven identity; mismatch, unproven round, or unpaired cumulative
  poisons to absent; a later matching round does not heal a mixed turn. 4 pinned
  fold tests (mismatch-poisons, unproven-poisons, never-heals, consistent-stays).
- buzz-acp/usage: pending_identity tri-state accumulator in UsageTracker; folded
  across in-flight notifications instead of last-update-wins. begin_turn() resets
  it; take() drains and stamps into the published TurnUsage. 3 pinned ACP
  behavior tests (A->B poisons, A->absent poisons, A->absent->A never heals).
  seed_zero_baseline now seeds last_total/last_cached_input/last_cache_write as
  Some(0): a freshly-spawned session has accumulated nothing, so the first turn's
  snapshot values are exact deltas, not discarded as 'no prior baseline'. 1 pinned
  first-turn round-trip test covering both cache categories.
- buzz-core/agent_turn_metric: PricingIdentity and cache-write fields.
- M2 migration (desktop/src-tauri): per-column guards matching M3 pattern —
  turn_cache_read_tokens and cumulative_cache_read_tokens checked and added
  independently so a crash between the two ALTERs leaves a self-repairable
  partial schema. Marker only commits after both columns are present. 2 new
  partial-schema repair tests: turn-only and cumulative-only shapes.
- M3 migration (desktop/src-tauri): adds turn_cache_write_tokens,
  cumulative_cache_write_tokens, pricing_authority, pricing_model,
  pricing_cache_class to agent_metric_index. Additive, idempotent,
  guarded by marker. Fresh-DB schema includes all M3 columns. 2 round-trip
  tests: positive full-row (nonzero cache-write + pricingIdentity) and
  omission/null case.
- tauriArchive.ts ReportedUsage: added cacheReadTokens, cacheWriteTokens,
  freshInputTokens as UsageField members to mirror the Rust struct field-for-field.
  1 pinned serde key-shape test: reported_usage_serializes_with_all_seven_camel_case_keys
  asserts exact camelCase key set so future Rust renames fail the gate.

P4a — aggregation layer (agent_usage.rs):

- Extended S-1 ladder to cache-read and cache-write via the same
  ladder_token path as input/output/total.
- derive_fresh_input: checked arithmetic, fail-closed — absent cache
  fields produce Unknown (not zero), overflow and subset > input both
  produce incomplete. Aggregated as a UsageField.
- D6 comparator: sort_value() = provider total when known, else
  input+output when both known, else None (unknown-last). Replaces the
  prior total-only comparator for both agent-level and model-level sort.
- Tests split into agent_usage_tests.rs (existing) and
  agent_usage_p4a_tests.rs (14 new pinned tests: cache ladder, fresh-input
  fail-closed cases, D6 comparator vector) to stay under the 1000-line ratchet.

Gate: cargo test -p buzz-agent (411/0), cargo test -p buzz-acp (682/0),
desktop src-tauri tests (2261/0, 14 ignored), just desktop-check (clean).

Co-authored-by: Will Pfleger <pfleger.will@gmail.com>
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch from 45ae876 to 0124f68 Compare August 5, 2026 22:31
…cumulation

input and output token accumulators used saturating_add throughout,
while the PR's own total and cache counters already used checked_add
with Unknown/Poisoned poison semantics (flagged by Wes's agent review).

Introduce TurnIOState (Unseen | Exact(n) | Poisoned) in types.rs:
- fold_round(n): overflow -> Poisoned, absence handled by caller
- merge_session(turn): Poisoned poisons permanently; Exact checked-add
- exact_value(): Some(n) only for Exact

Wire changes:
- accumulatedInputTokens / accumulatedOutputTokens: emitted only when
  Exact; omitted (never null, never u64::MAX) when Poisoned
- used: saturating fallback to 0 for display-only proxy (ACP dead code)

ACP changes:
- UsageUpdatePayload.accumulated_{input,output}_tokens: Option<u64>
  with #[serde(default)] for backward compat; None = publisher-poisoned
- absent field -> delta_reliable: false, null turn fields, null cumulative
- SessionState.last_{input,output}: Option<u64> to match provenance
- TurnUsage.cumulative_{input,output}_tokens: Option<u64>
- pool.rs: pass through directly (TokenCounts fields already Option<u64>)
- Box<UsageUpdatePayload> in GooseSessionUpdateVariant to satisfy
  clippy::large_enum_variant after struct grew with new doc+Option fields

Parser overflow fix (sum_usage: checked_add, SumUsageResult: Exact/Overflow):
- anthropic_input_tokens() returns Option<SumUsageResult>; Overflow sets
  LlmResponse.input_tokens_overflowed=true, None for input_tokens
- run loop: overflow flag -> TurnIOState::Poisoned, baseline frozen
- bool flag preferred over enum on LlmResponse: enum would ripple into
  ~20 existing test assertions on r.input_tokens == Some(...); flag
  confines the change to the two call sites that check it.

End-to-end regression: golden_transcripts drives real subprocess with
Anthropic-shaped response (input_tokens: u64::MAX, cache_read: 1) and
asserts accumulatedInputTokens absent from the usage_update notification.

Tests added (440+691 pass, 15 golden transcripts):
- TurnIOState unit tests: fold overflow, merge_session overflow, exact_value
- wire: omit input/output when poisoned; emit when exact
- ACP: absent input/output -> unreliable; goose payload unchanged;
  round-trip: poison mid-session -> subsequent turns stay unknown
- sum_usage overflow signal, parser overflow flag, parse-to-wire e2e

Co-authored-by: Will Pfleger <pfleger.will@gmail.com>
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
@wpfleger96
wpfleger96 force-pushed the duncan/agent-usage-backend branch from 0124f68 to 7e3762d Compare August 5, 2026 23:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants